import geopandas as gpd
import trimesh
import numpy as np
import os
from shapely.validation import make_valid
from shapely import is_valid
from shapely.geometry import box, Polygon, MultiPolygon
from shapely.ops import unary_union

# ---------------------- 核心工具函数 ----------------------
def create_contiguous_regions(geojson_path, num_regions=20):
    """创建连续无间隙的区域划分"""
    # 读取数据并重置索引，避免索引问题
    gdf = gpd.read_file(geojson_path, crs="EPSG:4326").reset_index(drop=True)

    # 修复几何并合并所有地块得到整体边界
    valid_geoms = []
    for idx, geom in enumerate(gdf.geometry):
        if geom is not None and not geom.is_empty:
            try:
                if not is_valid(geom):
                    fixed_geom = make_valid(geom)
                    valid_geoms.append(fixed_geom)
                else:
                    valid_geoms.append(geom)
            except Exception as e:
                print(f"【警告】跳过无效几何 {idx}：{str(e)}")
                continue

    if not valid_geoms:
        raise ValueError("没有有效的地块几何数据")

    # 计算整体边界
    union_geom = unary_union(valid_geoms)
    minx, miny, maxx, maxy = union_geom.bounds
    width = maxx - minx
    height = maxy - miny

    # 决定是按行还是按列划分
    if width > height:
        # 按宽度划分成垂直条带
        step = width / num_regions
        regions = []
        for i in range(num_regions):
            left = minx + i * step
            right = minx + (i + 1) * step if i < num_regions - 1 else maxx
            region_geom = box(left, miny, right, maxy)
            regions.append(region_geom)
    else:
        # 按高度划分成水平条带
        step = height / num_regions
        regions = []
        for i in range(num_regions):
            bottom = miny + i * step
            top = miny + (i + 1) * step if i < num_regions - 1 else maxy
            region_geom = box(minx, bottom, maxx, top)
            regions.append(region_geom)

    return gdf, regions

def get_parcels_in_region(gdf, region_geom):
    """获取区域内的所有地块（修复索引问题）"""
    # 创建区域GeoDataFrame，使用明确的索引
    region_gdf = gpd.GeoDataFrame(
        {'region_id': [0]},
        geometry=[region_geom],
        crs=gdf.crs,
        index=[0]
    )

    # 空间连接：找到与区域相交的地块，使用更安全的索引处理
    try:
        # 方法1：使用spatial_index进行查询（更稳定）
        spatial_index = gdf.sindex
        possible_matches_index = list(spatial_index.intersection(region_geom.bounds))
        possible_matches = gdf.iloc[possible_matches_index]
        precise_matches = possible_matches[possible_matches.intersects(region_geom)]
        return precise_matches

    except Exception as e:
        print(f"【警告】空间索引查询失败，使用备选方法：{str(e)}")
        # 方法2：使用sjoin但确保索引正确
        gdf_temp = gdf.copy().reset_index(drop=True)
        joined = gpd.sjoin(
            gdf_temp,
            region_gdf,
            how="inner",
            predicate="intersects"
        )
        # 移除连接产生的额外索引列
        joined = joined.drop(columns=['index_right'], errors='ignore')
        return joined

def create_parcel_mesh(geom, base_height=0.0, height=2.0):
    """创建单个地块的3D网格"""
    if geom is None or geom.is_empty:
        return None

    # 处理MultiPolygon
    if isinstance(geom, MultiPolygon):
        geoms = list(geom.geoms)
    elif isinstance(geom, Polygon):
        geoms = [geom]
    else:
        return None

    meshes = []
    for poly in geoms:
        # 获取外边界
        exterior = np.array(poly.exterior.coords)[:, :2]
        if len(exterior) < 4:
            continue

        # 创建底部
        bottom_verts = np.hstack([exterior, np.full((len(exterior), 1), base_height)])
        # 创建顶部
        top_verts = np.hstack([exterior, np.full((len(exterior), 1), base_height + height)])
        vertices = np.vstack([bottom_verts, top_verts])

        # 创建底部面
        bottom_faces = []
        for i in range(1, len(exterior) - 1):
            bottom_faces.append([0, i, i + 1])

        # 创建顶部面
        top_faces = []
        offset = len(exterior)
        for i in range(1, len(exterior) - 1):
            top_faces.append([offset, offset + i, offset + i + 1])

        # 创建侧面
        side_faces = []
        for i in range(len(exterior)):
            j = (i + 1) % len(exterior)
            side_faces.append([i, j, offset + j])
            side_faces.append([i, offset + j, offset + i])

        # 合并所有面
        if bottom_faces and top_faces and side_faces:
            faces = np.vstack([bottom_faces, top_faces, side_faces]).astype(np.int32)
            mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
            meshes.append(mesh)

    # 合并同一地块的多个部分
    if meshes:
        combined = trimesh.util.concatenate(meshes)
        # 设置颜色（每个地块颜色略有不同）
        color = np.random.randint(100, 200, 3)
        combined.visual.vertex_colors = np.append(color, 255)  # 添加alpha通道
        return combined
    return None

def export_region_fbx(region_idx, region_geom, parcels_gdf, output_dir, scale=100000):
    """导出单个区域的FBX文件，包含区域内所有地块"""
    scene = trimesh.Scene()

    # 计算区域中心和偏移，确保坐标合适
    centroid = region_geom.centroid
    region_minx, region_miny, _, _ = region_geom.bounds
    offset = np.array([region_minx, region_miny])

    # 添加区域内的所有地块
    for idx, row in parcels_gdf.iterrows():
        geom = row['geometry']
        try:
            # 转换坐标
            geom_coords = np.array(geom.centroid.coords)[0]
            translated_coords = (geom_coords - offset) * scale * 0.00001  # 调整缩放比例

            # 创建地块模型
            mesh = create_parcel_mesh(geom)
            if mesh is not None:
                # 移动到正确位置
                mesh.apply_translation([translated_coords[0], translated_coords[1], 0])

                # 设置地块名称
                parcel_id = str(row.get('MZGUID', idx))[:8]
                mesh.name = f"parcel_{parcel_id}"
                scene.add_geometry(mesh)
        except Exception as e:
            print(f"【警告】处理地块{idx}失败：{str(e)}")
            continue

    # 确保场景不为空
    if len(scene.geometry) == 0:
        print(f"【警告】区域{region_idx}没有有效地块，创建占位模型")
        placeholder = trimesh.creation.box(extents=[10, 10, 5])
        placeholder.visual.vertex_colors = [200, 100, 100, 255]  # 红色
        scene.add_geometry(placeholder)

    # 导出FBX
    file_name = f"land_region_{region_idx:02d}.fbx"
    output_path = os.path.join(output_dir, file_name)

    try:
        scene.export(file_type='fbx', file_obj=output_path)
        print(f"【成功】区域{region_idx}导出至：{output_path}，包含{len(scene.geometry)}个地块")
        return True
    except Exception as e:
        print(f"【错误】区域{region_idx}FBX导出失败：{str(e)}")
        # 尝试导出为OBJ
        obj_path = output_path.replace('.fbx', '.obj')
        try:
            scene.export(file_type='obj', file_obj=obj_path)
            print(f"【备选】区域{region_idx}导出为OBJ：{obj_path}")
            return True
        except:
            print(f"【失败】区域{region_idx}导出失败")
            return False

# ---------------------- 主函数 ----------------------
def export_20_contiguous_regions(geojson_path, output_dir):
    """将地块划分成20个连续无间隙区域并导出"""
    # 输入校验
    if not os.path.exists(geojson_path):
        print(f"【错误】输入GeoJSON不存在：{geojson_path}")
        return

    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"【提示】已创建输出目录：{output_dir}")

    try:
        # 创建20个连续无间隙区域
        print("【处理】正在创建20个连续无间隙区域...")
        gdf, regions = create_contiguous_regions(geojson_path, num_regions=20)

        # 为每个区域导出FBX
        for i, region in enumerate(regions, 1):
            print(f"\n【处理】区域{i}/20")
            # 获取区域内的所有地块（使用修复后的方法）
            parcels_in_region = get_parcels_in_region(gdf, region)
            print(f"【区域{i}】包含{len(parcels_in_region)}个地块")

            # 导出该区域
            export_region_fbx(
                region_idx=i,
                region_geom=region,
                parcels_gdf=parcels_in_region,
                output_dir=output_dir
            )

        print(f"\n【🎉 全部完成】导出目录：{output_dir}")
        print(f"共导出20个区域文件，每个区域包含其范围内的所有地块")

    except Exception as e:
        print(f"【错误】处理过程失败：{str(e)}")
        # 打印详细错误信息以便调试
        import traceback
        traceback.print_exc()
# ---------------------- 执行入口 ----------------------
if __name__ == "__main__":
    # 配置路径
    INPUT_GEOJSON = r"C:\Users\ma\Desktop\山河屯数据\林场信息.geojson"  # 替换为你的GeoJSON路径
    OUTPUT_DIR = r"C:\Users\ma\Desktop\山河屯数据\20_cubes_batches"       # 替换为你的输出目录

    # 导出20个地块批次
    export_20_contiguous_regions(
        geojson_path=INPUT_GEOJSON,
        output_dir=OUTPUT_DIR
    )
